"""
Utility functions for locELM package.

Includes error computation, visualization, and helper functions.
"""

import numpy as np
import matplotlib.pyplot as plt
from typing import Callable, Optional, Tuple
import time


def compute_errors(solution_func: Callable,
                  exact_func: Callable,
                  domain_bounds: list,
                  n_eval_points: list = None) -> dict:
    """
    Compute L-infinity and L2 errors.

    Parameters
    ----------
    solution_func : Callable
        Function that evaluates the computed solution
    exact_func : Callable
        Exact solution function
    domain_bounds : list
        Domain bounds [(a1,b1), (a2,b2), ...]
    n_eval_points : list
        Number of evaluation points per dimension

    Returns
    -------
    dict
        Dictionary with 'max_error' and 'rms_error'
    """
    ndim = len(domain_bounds)
    if n_eval_points is None:
        n_eval_points = [100] * ndim

    # Generate evaluation grid
    grids = []
    for (a, b), n in zip(domain_bounds, n_eval_points):
        grids.append(np.linspace(a, b, n))

    mesh = np.meshgrid(*grids, indexing='ij')
    eval_points = np.stack([m.flatten() for m in mesh], axis=1)

    # Compute solutions
    u_computed = solution_func(eval_points)
    u_exact = exact_func(eval_points)

    # Compute errors
    abs_error = np.abs(u_computed - u_exact)
    max_error = np.max(abs_error)
    rms_error = np.sqrt(np.mean(abs_error**2))

    return {
        'max_error': max_error,
        'rms_error': rms_error,
        'abs_error': abs_error,
        'eval_points': eval_points
    }


def plot_solution_1d(solution_func: Callable,
                    domain_bounds: Tuple[float, float],
                    exact_func: Optional[Callable] = None,
                    n_points: int = 200,
                    title: str = 'Solution',
                    save_path: Optional[str] = None):
    """
    Plot 1D solution.

    Parameters
    ----------
    solution_func : Callable
        Solution function
    domain_bounds : Tuple[float, float]
        Domain bounds (a, b)
    exact_func : Optional[Callable]
        Exact solution for comparison
    n_points : int
        Number of plot points
    title : str
        Plot title
    save_path : Optional[str]
        Path to save figure
    """
    a, b = domain_bounds
    x = np.linspace(a, b, n_points).reshape(-1, 1)
    u = solution_func(x)

    plt.figure(figsize=(10, 6))

    plt.subplot(1, 2, 1)
    plt.plot(x, u, 'b-', linewidth=2, label='Computed')
    if exact_func is not None:
        u_exact = exact_func(x)
        plt.plot(x, u_exact, 'r--', linewidth=2, label='Exact')
    plt.xlabel('x', fontsize=12)
    plt.ylabel('u', fontsize=12)
    plt.title(title, fontsize=14)
    plt.legend()
    plt.grid(True)

    if exact_func is not None:
        plt.subplot(1, 2, 2)
        error = np.abs(u - u_exact)
        plt.semilogy(x, error, 'g-', linewidth=2)
        plt.xlabel('x', fontsize=12)
        plt.ylabel('Absolute Error', fontsize=12)
        plt.title('Error Distribution', fontsize=14)
        plt.grid(True)

    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
    else:
        plt.show()


def plot_solution_2d(solution_func: Callable,
                    domain_bounds: list,
                    exact_func: Optional[Callable] = None,
                    n_points: list = None,
                    title: str = 'Solution',
                    save_path: Optional[str] = None):
    """
    Plot 2D solution.

    Parameters
    ----------
    solution_func : Callable
        Solution function
    domain_bounds : list
        Domain bounds [(a1,b1), (a2,b2)]
    exact_func : Optional[Callable]
        Exact solution for comparison
    n_points : list
        Number of plot points per dimension
    title : str
        Plot title
    save_path : Optional[str]
        Path to save figure
    """
    if n_points is None:
        n_points = [100, 100]

    (a1, b1), (a2, b2) = domain_bounds
    x = np.linspace(a1, b1, n_points[0])
    y = np.linspace(a2, b2, n_points[1])
    X, Y = np.meshgrid(x, y, indexing='ij')

    points = np.stack([X.flatten(), Y.flatten()], axis=1)
    u = solution_func(points).reshape(X.shape)

    if exact_func is not None:
        u_exact = exact_func(points).reshape(X.shape)
        fig, axes = plt.subplots(1, 3, figsize=(18, 5))
    else:
        fig, axes = plt.subplots(1, 1, figsize=(8, 6))
        axes = [axes]

    # Plot computed solution
    im = axes[0].contourf(X, Y, u, levels=50, cmap='viridis')
    axes[0].set_xlabel('x', fontsize=12)
    axes[0].set_ylabel('y', fontsize=12)
    axes[0].set_title(f'{title} (Computed)', fontsize=14)
    plt.colorbar(im, ax=axes[0])

    if exact_func is not None:
        # Plot exact solution
        im = axes[1].contourf(X, Y, u_exact, levels=50, cmap='viridis')
        axes[1].set_xlabel('x', fontsize=12)
        axes[1].set_ylabel('y', fontsize=12)
        axes[1].set_title(f'{title} (Exact)', fontsize=14)
        plt.colorbar(im, ax=axes[1])

        # Plot error
        error = np.abs(u - u_exact)
        im = axes[2].contourf(X, Y, error, levels=50, cmap='hot')
        axes[2].set_xlabel('x', fontsize=12)
        axes[2].set_ylabel('y', fontsize=12)
        axes[2].set_title('Absolute Error', fontsize=14)
        plt.colorbar(im, ax=axes[2])

    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
    else:
        plt.show()


class Timer:
    """Simple timer class for performance measurement."""

    def __init__(self):
        self.start_time = None
        self.elapsed_time = None

    def start(self):
        """Start timer."""
        self.start_time = time.time()

    def stop(self):
        """Stop timer and return elapsed time."""
        if self.start_time is None:
            raise RuntimeError("Timer not started")
        self.elapsed_time = time.time() - self.start_time
        return self.elapsed_time

    def __enter__(self):
        """Context manager entry."""
        self.start()
        return self

    def __exit__(self, *args):
        """Context manager exit."""
        self.stop()


def print_results(errors: dict, timing: float, method_name: str = "locELM"):
    """
    Print formatted results.

    Parameters
    ----------
    errors : dict
        Error dictionary from compute_errors
    timing : float
        Computation time in seconds
    method_name : str
        Name of method
    """
    print(f"\n{'='*60}")
    print(f"{method_name} Results")
    print(f"{'='*60}")
    print(f"Maximum Error (L∞): {errors['max_error']:.6e}")
    print(f"RMS Error (L2):     {errors['rms_error']:.6e}")
    print(f"Training Time:      {timing:.4f} seconds")
    print(f"{'='*60}\n")


def create_evaluation_function(networks, domain, output_dim=1):
    """
    Create a function that evaluates the solution at any points.

    Parameters
    ----------
    networks : MultiSubdomainNetwork
        Trained networks
    domain : Domain
        Domain decomposition
    output_dim : int
        Output dimension

    Returns
    -------
    Callable
        Function that evaluates solution at given points
    """
    def evaluate(points: np.ndarray) -> np.ndarray:
        """Evaluate solution at points."""
        points = np.atleast_2d(points)
        n_points = points.shape[0]
        result = np.zeros((n_points, output_dim))

        for i, point in enumerate(points):
            # Find which subdomain this point belongs to
            for subdomain_idx in range(domain.n_total):
                region = domain.get_subdomain_region(subdomain_idx)
                in_subdomain = True

                for dim, (a, b) in enumerate(region):
                    if point[dim] < a - 1e-10 or point[dim] > b + 1e-10:
                        in_subdomain = False
                        break

                if in_subdomain:
                    network = networks.get_network(subdomain_idx)
                    result[i] = network(point.reshape(1, -1)).flatten()
                    break

        return result.squeeze() if output_dim == 1 else result

    return evaluate
